#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdint.h>

#define RSA_KEY_SIZE 128 // For 1024-bit RSA key, adjust according to your key size

// Function to add PKCS#1 v1.5 padding
void rsa_pkcs1_padding(const uint8_t *data, size_t data_len, uint8_t *padded, size_t *padded_len) {
    if (data_len > (RSA_KEY_SIZE - 11)) {
        // Data too large to be padded
        *padded_len = 0;
        return;
    }

    // Set padding block
    memset(padded, 0x00, RSA_KEY_SIZE); // Clear the whole block
    padded[0] = 0x00;
    padded[1] = 0x02;
    padded[2] = 0xFF; // Random non-zero bytes (0xFF used here for simplicity)

    // Fill with padding bytes (0xFF), except the last byte
    size_t padding_len = RSA_KEY_SIZE - data_len - 3;
    memset(padded + 3, 0xFF, padding_len);
    padded[3 + padding_len] = 0x00;

    // Append the data
    memcpy(padded + 3 + padding_len + 1, data, data_len);

    *padded_len = RSA_KEY_SIZE;
}

// Function to remove PKCS#1 v1.5 padding
int rsa_pkcs1_unpadding(const uint8_t *padded, size_t padded_len, uint8_t *data, size_t *data_len) {
    if (padded_len != RSA_KEY_SIZE) {
        // Invalid length
        return -1;
    }

    if (padded[0] != 0x00 || padded[1] != 0x02) {
        // Invalid padding
        return -1;
    }

    size_t padding_len = 0;
    while (padded[2 + padding_len] == 0xFF) {
        padding_len++;
        if (2 + padding_len >= RSA_KEY_SIZE) {
            // Padding not found
            return -1;
        }
    }

    if (padded[2 + padding_len] != 0x00) {
        // Invalid padding byte
        return -1;
    }

    // Extract the data
    size_t data_start = 2 + padding_len + 1;
    size_t data_size = RSA_KEY_SIZE - data_start;
    if (data_size > padded_len) {
        // Data size out of bounds
        return -1;
    }

    memcpy(data, padded + data_start, data_size);
    *data_len = data_size;

    return 0;
}

int main() {
    uint8_t data[] = "Hello, World!";
    size_t data_len = strlen((char *)data);

    uint8_t padded[RSA_KEY_SIZE];
    size_t padded_len;

    rsa_pkcs1_padding(data, data_len, padded, &padded_len);

    printf("Padded data:\n");
    for (size_t i = 0; i < padded_len; i++) {
        printf("%02X ", padded[i]);
    }
    printf("\n");

    uint8_t unpadded[RSA_KEY_SIZE];
    size_t unpadded_len;

    if (rsa_pkcs1_unpadding(padded, padded_len, unpadded, &unpadded_len) == 0) {
        printf("Unpadded data: %.*s\n", (int)unpadded_len, unpadded);
    } else {
        printf("Failed to unpad data.\n");
    }

    return 0;
}
